from sko.tools import func_transformer
import numpy as np


class MyDeFunc():
  #初始化种群
  def __init__(self ,obj_func,size_pop, dim1, dim2, lb, ub,max_iter=100,
             prob_mut=0.3,CR=0.5,n_process=0,
             constraint_eq=tuple(),constraint_ueq=tuple()):
      assert size_pop % 2 == 0, 'size_pop must be even integer'
      self.size_pop = size_pop  # size of population
      self.max_iter = max_iter
      self.prob_mut = prob_mut  # probability of mutation
      self.n_dim1 = dim1
      self.n_dim2 = dim2
      self.CR=CR
      self.lb=lb
      self.ub=ub
      self.history_sum=[]
      self.test_sum=[]
      self.best_individual=[]
      # constraint:
      self.has_constraint = len(constraint_eq) > 0 or len(constraint_ueq) > 0
      self.constraint_eq = list(constraint_eq)  # a list of equal functions with ceq[i] = 0
      self.constraint_ueq = list(constraint_ueq)  # a list of unequal constraint functions with c[i] <= 0

      self.precision=None

      self.obj_func = obj_func
      self.population=None
      self.mut_population=None
      self.cro_population=None
      self.sel_population=None
      population=[]
      flag=0
      # 策略一
      # while True:
      #   target=[]
      #   for j in range(dim1):
      #     temp = np.random.uniform(lb[j], ub[j],  dim2)  # 生成一个1 x dim2的一维数组
      #     target.append(temp)
      #   target = np.array(target)
      #   if self.evaluate_fitness_individual(target) > 0 :
      #     population.append(target)  # 将这个二维数组添加到population列表中
      #     flag = flag+1
      #   if flag == size_pop:
      #      break
      # 策略二
      for _ in range(size_pop):
        target=[]
        for j in range(dim1):
            temp=np.random.uniform(lb[j],ub[j],dim2)
            target.append(temp)
        target=np.array(target)
        population.append(target)
      self.population = np.array(population)  # 将population列表转换为NumPy数组
      target = self.evaluate_fitness(self.population)

      # 找出最大适应度的索引
      max_index = np.argmax(target)
      # 使用索引取出最适应的并计算出它的目标函数
      self.best_individual=self.population[max_index]


  def evaluate_target_func(self):
      target=[]
      for i in range(len(self.population)):
        target.append(self.obj_func(self.population[i]))
      target = np.array(target)
      return target

  #求解适应度函数
  def evaluate_fitness(self,population):
      target=self.evaluate_target_func()
      if not self.has_constraint:
          target = target
      else:
        penalty_eq = np.array([np.sum(np.abs([c_i(x) for c_i in self.constraint_eq])) for x in population])
        penalty_ueq = np.array([np.sum(np.abs([np.maximum(0,c_i(x)) for c_i in self.constraint_ueq])) for x in population])
        fitness = target - penalty_eq - penalty_ueq
      return fitness
  #求解单个解的适应度函数
  def evaluate_fitness_individual(self,individual):
      target=self.obj_func(individual)
      if not self.has_constraint:
          target = target
      else:
        penalty_eq = np.array(np.sum(np.abs([c_i(individual) for c_i in self.constraint_eq])))
        penalty_ueq = np.array(np.sum(np.abs([np.maximum(0,c_i(individual)) for c_i in self.constraint_ueq])) )
        fitness = target - 1e3*penalty_eq - 1e3*penalty_ueq
      return fitness
  # 变异操作
  def mutate(self):
      self.mut_population=np.full_like(self.population,0)
      for i in range(len(self.population)):
          r_idx=None
          # 变异策略1
          # while True:
          #   r_idx = np.random.choice(len(self.population), 3, replace=False)
          #   if len(set(r_idx)) == 3 and r_idx[0] != i and r_idx[1] != i and r_idx[2] != i:
          #       break
          # r1, r2, r3 = r_idx
          # self.mut_population[i] = self.population[r1] + self.prob_mut * (self.population[r2] - self.population[r3])
          # 变异策略二
          while True:
            r_idx = np.random.choice(len(self.population), 2, replace=False)
            if len(set(r_idx)) == 2 and r_idx[0] != i and r_idx[1] != i:
                break
          r1, r2 = r_idx
          self.mut_population[i] = self.best_individual + self.prob_mut*(self.population[r1] - self.population[r2])
          
      # 边界值处理

      # 策略一：超出边界值随机处理
      mask=[]
      for i in range(self.size_pop):
        target=[]
        for j in range(self.n_dim1):
          temp = np.random.uniform(self.lb[j], self.ub[j],  self.n_dim2)  # 生成一个1 x dim2的二维数组
          target.append(temp)
        mask.append(np.array(target))    
      mask = np.array((mask))
      for i in range(len(self.mut_population)):
         for j in range(len(self.mut_population[0])):
            for k in range(len(self.mut_population[0][0])):            
              if self.mut_population[i][j][k] < self.lb[j]:
                self.mut_population[i][j][k]=mask[i][j][k]
              if self.mut_population[i][j][k] > self.ub[j]:
                self.mut_population[i][j][k]=mask[i][j][k]
      # 策略二 超出边界取边界值
      # for i in range(len(self.mut_population)):
      #   for j in range(len(self.mut_population[0])):
      #       for k in range(len(self.mut_population[0][0])):            
      #         if self.mut_population[i][j][k] < self.lb[j]:
      #           self.mut_population[i][j][k]=self.lb[j]
      #         if self.mut_population[i][j][k] > self.ub[j]:
      #           self.mut_population[i][j][k]=self.ub[j]
      return self.mut_population

  # 交叉操作
  def crossover(self):
      self.cro_population = np.copy(self.population)
      for i in range(len(self.population)):
          j = np.random.randint(0, len(self.population[0]))
          for k in range(len(self.population[0])):
              if np.random.rand() < self.CR or k == j:
                  self.cro_population[i][k] = self.mut_population[i][k]
      return self.cro_population


  def select(self):
      self.sel_population = np.copy(self.population)
      '''
      greedy selection
      '''
      fitness_raw = self.evaluate_fitness(self.population)
      fitness_cross = self.evaluate_fitness(self.cro_population)
      for i in range(len(self.population)):
          if fitness_raw[i] > fitness_cross[i]:
            self.sel_population[i] = self.cro_population[i]
      target = self.evaluate_fitness(self.sel_population)

      # 找出最大适应度的索引
      max_index = np.argmax(target)

      # 使用索引取出最适应的并计算出它的目标函数

      temp=self.obj_func(self.sel_population[max_index])
      self.history_sum.append(temp)
      # self.test_sum.append(self.obj_func(self.sel_population[0]))
      self.best_individual=self.sel_population[max_index]
      self.precision=temp-target[max_index]
      return self.sel_population

  # 【10个种群个体，【每个个体n个变量，【每个变量分为24个时间段】】】
  def run(self):
      for i in range(self.max_iter):
          self.mutate()
          self.crossover()
          self.population = self.select()
          print(f'第{i}次迭代\n')
          # if self.precision < 100:
          #    break
      return self